import numpy as np
from torch.nn import functional as F
import torch
from model.fusions.fusions import MFB
import torch.nn as nn

class ConcatELUAttn(nn.Module):
#<<<<<<< HEAD

#=======
#>>>>>>> 22b53cee88948f62283b7e04f4292db8c81c32b9
    def __init__(self, module_dim=768):
        super(ConcatELUAttn, self).__init__()
        self.module_dim = module_dim                         #输入模型维度

        self.q_proj = nn.Linear(module_dim, module_dim, bias=False)     #定义的几个层结构
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.cat = nn.Linear(2*module_dim, module_dim)
        self.attn = nn.Linear(module_dim, 1)

        self.activation = nn.ELU()
        self.dropout = nn.Dropout(0.15)

    def forward(self, question_rep, visual_feat):   #前向传播时传入两个张量  问题特征和视频特征
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        q_proj = self.q_proj(question_rep)
        v_proj = self.v_proj(visual_feat)
        q_proj = q_proj.unsqueeze(1)

        v_q_cat = torch.cat((v_proj, q_proj.expand_as(v_proj) * v_proj), dim=-1)
        v_q_cat = self.cat(v_q_cat)
        v_q_cat = self.activation(v_q_cat)

        attn = self.attn(v_q_cat)  # (bz, k, 1)
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)#用于对视觉特征进行加权

        v_distill = (attn * visual_feat).sum(1)    #注意力权重对视觉特征进行加权求和操作

        return v_distill

class MFBAttn(nn.Module):                
    def __init__(self, module_dim=768):
        super(MFBAttn, self).__init__()
        self.module_dim = module_dim

        self.q_proj = nn.Linear(module_dim, module_dim, bias=False) 
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.cat = MFB([module_dim, module_dim], module_dim, mm_dim=module_dim, factor=2)
        self.attn = nn.Linear(module_dim, 1)

        self.activation = nn.Tanh()
        self.dropout = nn.Dropout(0.15)

    def forward(self, question_rep, visual_feat):  #输入包括问题表示(question_rep)和视觉特征(visual_feat)，输出为视觉蒸馏表示(v_distill)
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        q_proj = self.q_proj(question_rep) #通过线性变换对问题表示(question_rep)和视觉特征(visual_feat)进行投影，得到q_proj和v_proj
        v_proj = self.v_proj(visual_feat)
        q_proj = q_proj.unsqueeze(1)      # 在第1个维度插入新的维度

        v_q_cat = self.cat(v_proj, q_proj.expand_as(v_proj) * v_proj)
        #通过将问题表示(q_proj)进行扩展并与视觉特征(v_proj)相乘，并调用cat函数（MFB），得到v_q_cat。v_q_cat是对问题表示和视觉特征进行拼接和交互的结果。

        attn = self.attn(v_q_cat)  # (bz, k, 1)    #将v_q_cat经过线性变换得到注意力分数attn，然后通过softmax函数进行归一化，得到具有权重的注意力分布。
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)

        v_distill = (attn * visual_feat).sum(1)      #将注意力分布与视觉特征相乘，并对所有特征求和，得到视觉蒸馏表示(v_distill)。

        return v_distill

class MFBSimpleAttn(nn.Module):   #和上一个基本一样，不一样的地方标注出来了
    def __init__(self, module_dim=768):
        super(MFBAttn, self).__init__()
        self.module_dim = module_dim

        self.q_proj = nn.Linear(module_dim, module_dim, bias=False)
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.cat = MFB([2 * module_dim, module_dim * 2], module_dim, mm_dim=module_dim, factor=2)
        #输入维度变化
        self.attn = nn.Linear(module_dim, 1)

        self.activation = nn.Tanh()
        self.dropout = nn.Dropout(0.15)

    def forward(self, question_rep, visual_feat):
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        q_proj = self.q_proj(question_rep)
        v_proj = self.v_proj(visual_feat)
        q_proj = q_proj.unsqueeze(1)

        v_q_cat = self.cat(v_proj, q_proj.expand_as(v_proj))
        #有点儿疑惑？
        attn = self.attn(v_q_cat)  # (bz, k, 1)
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)

        v_distill = (attn * visual_feat).sum(1)

        return v_distill

class SimpleConcatELUAttn(nn.Module):#该函数接收一个表示问题的张量 question_rep 和一个表示视觉特征的张量 visual_feat，并返回一个表示视觉蒸馏结果的张量 v_distill。
    def __init__(self, module_dim=768):
        super(SimpleConcatELUAttn, self).__init__()
        self.module_dim = module_dim

        self.q_proj = nn.Linear(module_dim, module_dim, bias=False)
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.cat = nn.Linear(2*module_dim, module_dim)
        self.attn = nn.Linear(module_dim, 1)

        self.activation = nn.ELU()
        self.dropout = nn.Dropout(0.15)

    def forward(self, question_rep, visual_feat):
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        q_proj = self.q_proj(question_rep)#通过线性层 q_proj 对问题表示 question_rep 进行线性变换，得到 q_proj。
        v_proj = self.v_proj(visual_feat)#通过线性层 v_proj 对视觉特征表示 visual_feat 进行线性变换，得到 v_proj。
        q_proj = q_proj.unsqueeze(1)
     #q_proj.unsqueeze(1) 是将 q_proj 张量在第一维上增加一个维度。将 question_rep 的形状从 (batch_size, module_dim) 
        # 转换为 (batch_size, 1, module_dim)，使其在第二个维度上与 visual_feat 的第二、第三个维度保持一致。这样就可以在后续
        # 的拼接操作中，将两者按照最后一个维度进行拼接，得到形状为 (batch_size, num_of_clips, 2*module_dim) 的张量。
        v_q_cat = torch.cat((v_proj, q_proj.expand_as(v_proj)), dim=-1)
        #并将两者在最后一个维度上进行拼接，得到 v_q_cat
        v_q_cat = self.cat(v_q_cat)    #通过线性层 cat 对 v_q_cat 进行线性变换，并应用 ELU 激活函数。
        v_q_cat = self.activation(v_q_cat)

        attn = self.attn(v_q_cat)  # (bz, k, 1)#通过线性层 attn 对 v_q_cat 进行线性变换，得到注意力权重 attn。
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)使用 softmax 函数对 attn 进行归一化，得到概率分布。

        v_distill = (attn * visual_feat).sum(1)#将视觉特征 visual_feat 与注意力权重 attn 相乘并在第二维上求和，得到视觉蒸馏结果 v_distill。

        return v_distill

class ContextSelfAttn(nn.Module):#该模型接收一个表示视觉特征的张量 visual_feat，经过线性变换、注意力计算和汇总等操作产生一个表示视觉蒸馏结果的张量 v_distill。
    def __init__(self, module_dim=768):
        super(ContextSelfAttn, self).__init__()
        self.module_dim = module_dim
        self.v_proj = nn.Linear(module_dim, module_dim, bias=False)

        self.attn = nn.Linear(module_dim, 1)  #创建了线性层 attn 用于计算注意力分数
        self.activation = nn.ELU()
        self.dropout = nn.Dropout(0.15)

    def forward(self, visual_feat):
        """
        Args:
            question_rep: [Tensor] (batch_size, module_dim)
            visual_feat: [Tensor] (batch_size, num_of_clips, module_dim)
        return:
            visual_distill representation [Tensor] (batch_size, module_dim)
        """
        visual_feat = self.dropout(visual_feat) # TODO
        v_proj = self.v_proj(visual_feat)
        v_q_cat = self.activation(v_proj)

        attn = self.attn(v_q_cat)  # (bz, k, 1)#线性层 attn 对 v_q_cat 进行线性变换，得到注意力权重 attn。
        attn = F.softmax(attn, dim=1)  # (bz, k, 1)#使用 softmax 函数对 attn 进行归一化，得到概率分布。

        v_distill = (attn * visual_feat).sum(1)#将概率分布 attn 与视觉特征 visual_feat 相乘，并在第二维上求和，得到视觉蒸馏结果 v_distill。

        return v_distill   #模型利用自注意力机制，对每个视觉特征进行注意力计算，从而聚焦于重要的特征表示，并生成一个视觉蒸馏结果。

class SimpleOutputUnitOpenEnded(nn.Module): # concat->Linear classifier
    def __init__(self, module_dim=512, num_answers=1000):#num_answers 参数指定了分类器的输出类别数目
        super(SimpleOutputUnitOpenEnded, self).__init__()
       # 该模型接收一个表示问题嵌入的张量
       # question_embedding和一个表示视觉嵌入的张量visual_embedding，通过线性分类器将它们进行拼接并输出最终的结果张量out。
        self.question_proj = nn.Linear(module_dim, module_dim)

        self.classifier = nn.Sequential(nn.Dropout(0.15),
                                        nn.Linear(module_dim * 2, module_dim),
                                        nn.ELU(),
                                        nn.BatchNorm1d(module_dim),
                                        nn.Dropout(0.15),
                                        nn.Linear(module_dim, num_answers))
  #创建了一个包含多个层的顺序模型 classifier，其中包括一个 dropout 层、两个线性层、一个 ELU 激活函数、一个归一化层和一个 dropout 层。
    def forward(self, question_embedding, visual_embedding):
        question_embedding = self.question_proj(question_embedding)
        out = torch.cat([visual_embedding, question_embedding], 1)
        #使用 torch.cat 函数将视觉嵌入和问题嵌入在第二维上进行拼接，得到一个维度为 module_dim * 2 的张量 out。
        out = self.classifier(out)
   #将 out 输入到顺序模型 classifier 中，经过一系列的层操作后，得到最终的结果张量 out。
        return out

class GateOutputUnitOpenEnded(nn.Module): # concat->Linear classifier，和上一个大体相同，不同处标注出来
    def __init__(self, module_dim=512, num_answers=1000):
        super(GateOutputUnitOpenEnded, self).__init__()

        self.question_proj = nn.Linear(module_dim, module_dim)

        self.classifier = nn.Sequential(nn.Dropout(0.15),
                                        nn.Linear(module_dim * 2, module_dim),
                                        nn.ELU(),
                                        nn.BatchNorm1d(module_dim),
                                        nn.Dropout(0.15),
                                        nn.Linear(module_dim, num_answers))
        self.gate = nn.Linear(module_dim * 2, module_dim * 2)
           #还创建了一个线性层 gate，用于计算门控机制的权重。
    def forward(self, question_embedding, visual_embedding):
        question_embedding = self.question_proj(question_embedding)
        out = torch.cat([visual_embedding, question_embedding], 1)
        gate = self.gate(out)   #将 out 输入到线性层 gate 中，得到门控机制的权重张量 gate。
        out = gate * out        #将 gate 与 out 相乘，对视觉嵌入和问题嵌入进行加权融合操作。
        out = self.classifier(out)

        return out

